{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47ed5db3-bda5-49e1-bf75-23e08c9a3af0",
   "metadata": {},
   "source": [
    "# How to pass private state\n",
    "\n",
    "Oftentimes, you may want nodes to be able to pass state to each other that should NOT be part of the main schema of the graph. This is often useful because there may be information that is not needed as input/output (and therefore doesn't really make sense to have in the main schema) but is needed as part of the intermediate working logic.\n",
    "\n",
    "Let's take a look at an example below. In this example, we will create a RAG pipeline that:\n",
    "1. Takes in a user question\n",
    "2. Uses an LLM to generate a search query\n",
    "3. Retrieves documents for that generated query\n",
    "4. Generates a final answer based on those documents\n",
    "\n",
    "We will have a separate node for each step. We will only have the `question` and `answer` on the overall state. However, we will need separate states for the `search_query` and the `documents` - we will pass these as private state keys by defining an `input` annotation on each relevant node.\n",
    "\n",
    "Let's look at an example!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3114c3ad-0ade-47ba-9488-53d6f7671578",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  question: 'How are you?',\n",
      "  answer: 'How are you? rephrased as a query!\\n\\nsome random document\\n\\nHow are you?'\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import { Annotation, StateGraph } from \"@langchain/langgraph\";\n",
    "\n",
    "\n",
    "// The overall state of the graph\n",
    "const OverallStateAnnotation = Annotation.Root({\n",
    "  question: Annotation<string>,\n",
    "  answer: Annotation<string>,\n",
    "});\n",
    "\n",
    "// This is what the node that generates the query will return\n",
    "const QueryOutputAnnotation = Annotation.Root({\n",
    "  query: Annotation<string>,\n",
    "});\n",
    "\n",
    "// This is what the node that retrieves the documents will return\n",
    "const DocumentOutputAnnotation = Annotation.Root({\n",
    "  docs: Annotation<string[]>,\n",
    "});\n",
    "\n",
    "// This is what the node that retrieves the documents will return\n",
    "const GenerateOutputAnnotation = Annotation.Root({\n",
    "  ...OverallStateAnnotation.spec,\n",
    "  ...DocumentOutputAnnotation.spec\n",
    "});\n",
    "\n",
    "// Node to generate query\n",
    "const generateQuery = async (state: typeof OverallStateAnnotation.State) => {\n",
    "  // Replace this with real logic\n",
    "  return {\n",
    "    query: state.question + \" rephrased as a query!\",\n",
    "  };\n",
    "};\n",
    "\n",
    "// Node to retrieve documents\n",
    "const retrieveDocuments = async (state: typeof QueryOutputAnnotation.State) => {\n",
    "  // Replace this with real logic\n",
    "  return {\n",
    "    docs: [state.query, \"some random document\"],\n",
    "  };\n",
    "};\n",
    "\n",
    "// Node to generate answer\n",
    "const generate = async (state: typeof GenerateOutputAnnotation.State) => {\n",
    "  return {\n",
    "    answer: state.docs.concat([state.question]).join(\"\\n\\n\"),\n",
    "  };\n",
    "};\n",
    "\n",
    "const graph = new StateGraph(OverallStateAnnotation)\n",
    "  .addNode(\"generate_query\", generateQuery)\n",
    "  .addNode(\"retrieve_documents\", retrieveDocuments, { input: QueryOutputAnnotation })\n",
    "  .addNode(\"generate\", generate, { input: GenerateOutputAnnotation })\n",
    "  .addEdge(\"__start__\", \"generate_query\")\n",
    "  .addEdge(\"generate_query\", \"retrieve_documents\")\n",
    "  .addEdge(\"retrieve_documents\", \"generate\")\n",
    "  .compile();\n",
    "\n",
    "await graph.invoke({\n",
    "  question: \"How are you?\",\n",
    "});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30cc74ea",
   "metadata": {},
   "source": [
    "Above, the original `question` value in the input has been preserved, but that the `generate_query` node rephrased it, the `retrieve_documents` node added `\"some random document\"`, and finally the `generate` node combined the `docs` in the state with the original question to create an `answer`. The intermediate steps populated by the `input` annotations passed to the individual nodes are not present in the final output."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TypeScript",
   "language": "typescript",
   "name": "tslab"
  },
  "language_info": {
   "codemirror_mode": {
    "mode": "typescript",
    "name": "javascript",
    "typescript": true
   },
   "file_extension": ".ts",
   "mimetype": "text/typescript",
   "name": "typescript",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
